import os
import numpy as np
from tqdm import tqdm
import pandas as pd
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei'] # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False # 用来正常显示负号

def plot_results(preds:np.ndarray,trues:np.ndarray,data_name:str,save_root:str):
    ylen = preds.shape[1]
    
    def save_fig(p:np.ndarray,t:np.ndarray,title,save_to):
        fig = plt.figure(figsize=(12,5),dpi=300)
        ax = fig.add_subplot()
        ax.set_xlim(0,ylen)
        ax.plot(p.sum(axis=1),'r',label="Pred")
        ax.plot(t.sum(axis=1),'b',label="True")
        ax.legend()
        ax.set_title(title)
        fig.savefig(save_to)
    
    if data_name.endswith("_TPLV"):
        heads = "CS1#c,CS10#c,CS2#c,CS3#c,CS4#c,CS5#c,CS6#c,CS7#c,CS8#c,CS9#c,edge10_1#c,edge10_6#c,edge10_7#c,edge11_12#c,edge11_3#c,edge11_8#c,edge12_11#c,edge12_13#c,edge12_17#c,edge12_35#c,edge13_12#c,edge13_14#c,edge13_21#c,edge13_8#c,edge14_13#c,edge14_15#c,edge14_17#c,edge15_14#c,edge15_16#c,edge15_20#c,edge16_15#c,edge16_17#c,edge16_19#c,edge16_31#c,edge17_12#c,edge17_14#c,edge17_16#c,edge17_32#c,edge17_35#c,edge18_19#c,edge18_29#c,edge18_37#c,edge19_16#c,edge19_18#c,edge19_20#c,edge19_25#c,edge1_10#c,edge1_2#c,edge1_5#c,edge20_15#c,edge20_19#c,edge20_21#c,edge20_25#c,edge21_13#c,edge21_20#c,edge21_22#c,edge21_24#c,edge22_21#c,edge22_23#c,edge22_7#c,edge23_22#c,edge23_24#c,edge23_28#c,edge23_6#c,edge24_21#c,edge24_23#c,edge24_25#c,edge24_27#c,edge25_19#c,edge25_20#c,edge25_24#c,edge25_26#c,edge25_29#c,edge26_25#c,edge26_27#c,edge26_30#c,edge27_24#c,edge27_26#c,edge27_28#c,edge28_23#c,edge28_27#c,edge28_5#c,edge29_18#c,edge29_25#c,edge29_30#c,edge29_37#c,edge2_1#c,edge2_3#c,edge2_8#c,edge30_26#c,edge30_29#c,edge30_9#c,edge31_16#c,edge31_32#c,edge31_37#c,edge31_9#c,edge32_17#c,edge32_31#c,edge32_33#c,edge32_34#c,edge33_32#c,edge33_34#c,edge33_35#c,edge33_4#c,edge34_32#c,edge34_33#c,edge34_36#c,edge35_12#c,edge35_17#c,edge35_33#c,edge35_4#c,edge36_34#c,edge36_9#c,edge37_18#c,edge37_29#c,edge37_31#c,edge3_11#c,edge3_2#c,edge3_4#c,edge4_3#c,edge4_33#c,edge4_35#c,edge5_1#c,edge5_28#c,edge5_6#c,edge6_10#c,edge6_23#c,edge6_5#c,edge6_7#c,edge7_10#c,edge7_22#c,edge7_6#c,edge7_8#c,edge8_11#c,edge8_13#c,edge8_2#c,edge8_7#c,edge9_30#c,edge9_31#c,edge9_36#c,edge10_1#d,edge10_6#d,edge10_7#d,edge11_12#d,edge11_3#d,edge11_8#d,edge12_11#d,edge12_13#d,edge12_17#d,edge12_35#d,edge13_12#d,edge13_14#d,edge13_21#d,edge13_8#d,edge14_13#d,edge14_15#d,edge14_17#d,edge15_14#d,edge15_16#d,edge15_20#d,edge16_15#d,edge16_17#d,edge16_19#d,edge16_31#d,edge17_12#d,edge17_14#d,edge17_16#d,edge17_32#d,edge17_35#d,edge18_19#d,edge18_29#d,edge18_37#d,edge19_16#d,edge19_18#d,edge19_20#d,edge19_25#d,edge1_10#d,edge1_2#d,edge1_5#d,edge20_15#d,edge20_19#d,edge20_21#d,edge20_25#d,edge21_13#d,edge21_20#d,edge21_22#d,edge21_24#d,edge22_21#d,edge22_23#d,edge22_7#d,edge23_22#d,edge23_24#d,edge23_28#d,edge23_6#d,edge24_21#d,edge24_23#d,edge24_25#d,edge24_27#d,edge25_19#d,edge25_20#d,edge25_24#d,edge25_26#d,edge25_29#d,edge26_25#d,edge26_27#d,edge26_30#d,edge27_24#d,edge27_26#d,edge27_28#d,edge28_23#d,edge28_27#d,edge28_5#d,edge29_18#d,edge29_25#d,edge29_30#d,edge29_37#d,edge2_1#d,edge2_3#d,edge2_8#d,edge30_26#d,edge30_29#d,edge30_9#d,edge31_16#d,edge31_32#d,edge31_37#d,edge31_9#d,edge32_17#d,edge32_31#d,edge32_33#d,edge32_34#d,edge33_32#d,edge33_34#d,edge33_35#d,edge33_4#d,edge34_32#d,edge34_33#d,edge34_36#d,edge35_12#d,edge35_17#d,edge35_33#d,edge35_4#d,edge36_34#d,edge36_9#d,edge37_18#d,edge37_29#d,edge37_31#d,edge3_11#d,edge3_2#d,edge3_4#d,edge4_3#d,edge4_33#d,edge4_35#d,edge5_1#d,edge5_28#d,edge5_6#d,edge6_10#d,edge6_23#d,edge6_5#d,edge6_7#d,edge7_10#d,edge7_22#d,edge7_6#d,edge7_8#d,edge8_11#d,edge8_13#d,edge8_2#d,edge8_7#d,edge9_30#d,edge9_31#d,edge9_36#d,edge10_1#v2g,edge10_6#v2g,edge10_7#v2g,edge11_12#v2g,edge11_3#v2g,edge11_8#v2g,edge12_11#v2g,edge12_13#v2g,edge12_17#v2g,edge12_35#v2g,edge13_12#v2g,edge13_14#v2g,edge13_21#v2g,edge13_8#v2g,edge14_13#v2g,edge14_15#v2g,edge14_17#v2g,edge15_14#v2g,edge15_16#v2g,edge15_20#v2g,edge16_15#v2g,edge16_17#v2g,edge16_19#v2g,edge16_31#v2g,edge17_12#v2g,edge17_14#v2g,edge17_16#v2g,edge17_32#v2g,edge17_35#v2g,edge18_19#v2g,edge18_29#v2g,edge18_37#v2g,edge19_16#v2g,edge19_18#v2g,edge19_20#v2g,edge19_25#v2g,edge1_10#v2g,edge1_2#v2g,edge1_5#v2g,edge20_15#v2g,edge20_19#v2g,edge20_21#v2g,edge20_25#v2g,edge21_13#v2g,edge21_20#v2g,edge21_22#v2g,edge21_24#v2g,edge22_21#v2g,edge22_23#v2g,edge22_7#v2g,edge23_22#v2g,edge23_24#v2g,edge23_28#v2g,edge23_6#v2g,edge24_21#v2g,edge24_23#v2g,edge24_25#v2g,edge24_27#v2g,edge25_19#v2g,edge25_20#v2g,edge25_24#v2g,edge25_26#v2g,edge25_29#v2g,edge26_25#v2g,edge26_27#v2g,edge26_30#v2g,edge27_24#v2g,edge27_26#v2g,edge27_28#v2g,edge28_23#v2g,edge28_27#v2g,edge28_5#v2g,edge29_18#v2g,edge29_25#v2g,edge29_30#v2g,edge29_37#v2g,edge2_1#v2g,edge2_3#v2g,edge2_8#v2g,edge30_26#v2g,edge30_29#v2g,edge30_9#v2g,edge31_16#v2g,edge31_32#v2g,edge31_37#v2g,edge31_9#v2g,edge32_17#v2g,edge32_31#v2g,edge32_33#v2g,edge32_34#v2g,edge33_32#v2g,edge33_34#v2g,edge33_35#v2g,edge33_4#v2g,edge34_32#v2g,edge34_33#v2g,edge34_36#v2g,edge35_12#v2g,edge35_17#v2g,edge35_33#v2g,edge35_4#v2g,edge36_34#v2g,edge36_9#v2g,edge37_18#v2g,edge37_29#v2g,edge37_31#v2g,edge3_11#v2g,edge3_2#v2g,edge3_4#v2g,edge4_3#v2g,edge4_33#v2g,edge4_35#v2g,edge5_1#v2g,edge5_28#v2g,edge5_6#v2g,edge6_10#v2g,edge6_23#v2g,edge6_5#v2g,edge6_7#v2g,edge7_10#v2g,edge7_22#v2g,edge7_6#v2g,edge7_8#v2g,edge8_11#v2g,edge8_13#v2g,edge8_2#v2g,edge8_7#v2g,edge9_30#v2g,edge9_31#v2g,edge9_36#v2g".split(",")
        cs_c = slice(0,10)
        pk_c = slice(10,140)
        pk_d = slice(140,270)
        pk_v2g = slice(270,400)
    else:
        heads = "CS1,CS10,CS2,CS3,CS4,CS5,CS6,CS7,CS8,CS9,edge10_1,edge10_6,edge10_7,edge11_12,edge11_3,edge11_8,edge12_11,edge12_13,edge12_17,edge12_35,edge13_12,edge13_14,edge13_21,edge13_8,edge14_13,edge14_15,edge14_17,edge15_14,edge15_16,edge15_20,edge16_15,edge16_17,edge16_19,edge16_31,edge17_12,edge17_14,edge17_16,edge17_32,edge17_35,edge18_19,edge18_29,edge18_37,edge19_16,edge19_18,edge19_20,edge19_25,edge1_10,edge1_2,edge1_5,edge20_15,edge20_19,edge20_21,edge20_25,edge21_13,edge21_20,edge21_22,edge21_24,edge22_21,edge22_23,edge22_7,edge23_22,edge23_24,edge23_28,edge23_6,edge24_21,edge24_23,edge24_25,edge24_27,edge25_19,edge25_20,edge25_24,edge25_26,edge25_29,edge26_25,edge26_27,edge26_30,edge27_24,edge27_26,edge27_28,edge28_23,edge28_27,edge28_5,edge29_18,edge29_25,edge29_30,edge29_37,edge2_1,edge2_3,edge2_8,edge30_26,edge30_29,edge30_9,edge31_16,edge31_32,edge31_37,edge31_9,edge32_17,edge32_31,edge32_33,edge32_34,edge33_32,edge33_34,edge33_35,edge33_4,edge34_32,edge34_33,edge34_36,edge35_12,edge35_17,edge35_33,edge35_4,edge36_34,edge36_9,edge37_18,edge37_29,edge37_31,edge3_11,edge3_2,edge3_4,edge4_3,edge4_33,edge4_35,edge5_1,edge5_28,edge5_6,edge6_10,edge6_23,edge6_5,edge6_7,edge7_10,edge7_22,edge7_6,edge7_8,edge8_11,edge8_13,edge8_2,edge8_7,edge9_30,edge9_31,edge9_36".split(',')
        cs_c = slice(0,10)
        pk_c = slice(10,140)
        pk_d = slice(0,0)
        pk_v2g = slice(0,0)
    
    os.makedirs(save_root, exist_ok=True)
    for i in tqdm(range(preds.shape[0]), mininterval=1, desc=f'  - (Plotting) ', leave=False):
        this_root = f"{save_root}/{i}"
        os.makedirs(this_root, exist_ok=True)
        pred=preds[i];true=trues[i]
        save_fig(pred[:,cs_c],true[:,cs_c],"所有快充站的总负荷(kW)",f"{this_root}/load_cs.png")
        save_fig(pred[:,pk_c],true[:,pk_c],"所有路边慢充的总充电负荷(kW)",f"{this_root}/load_pk_c.png")
        if pk_d.start<pk_d.stop:
            save_fig(pred[:,pk_d],true[:,pk_d],"所有路边慢充的总放电功率(kW)",f"{this_root}/load_pk_d.png")
            save_fig(np.concatenate((pred[:,cs_c],pred[:,pk_c]),axis=1),
                        np.concatenate((true[:,cs_c],true[:,pk_c]),axis=1),"所有充电点的净负荷(kW)",f"{this_root}/load_all.png")
        else:
            save_fig(pred,true,"所有充电点的净负荷(kW)",f"{this_root}/load_all.png")
        if pk_v2g.start<pk_v2g.stop:
            save_fig(pred[:,pk_v2g],true[:,pk_v2g],"所有路边慢充的最大放电功率(kW)",f"{this_root}/load_pk_v2g.png")
        pd.DataFrame(data=pred,columns=heads).to_csv(f"{this_root}/preds.csv")
        pd.DataFrame(data=true,columns=heads).to_csv(f"{this_root}/trues.csv")